import datetime
import inspect
import itertools
import os
import pickle
import random
import traceback
import bpnn

from api import Commander
from api import commands
from api import gameinfo
from api.vector2 import Vector2

# NOTES
#
# (0,0) is top left

class TZCommander(Commander):
  # Tactics
  NONE = "None"
  GO_ENEMY_FLAG = "GoEnemyFlag"
  GO_OUR_FLAG = "GoOurFlag"
  INTERCEPT_FLAG = "InterceptFlag"
  GO_GOAL = "GoGoal"
  GO_ENEMY_GOAL = "GoEnemyGoal"
  PATROL = "Patrol"
  DEFENSE_PATROL = "DefensePatrol"
  ATTACK_BOT = "AttackBot"
  GO_ENEMY_FLAG_VIA_FLANK = "GoEnemyFlagViaFlank"

  DEAD = "Dead"
  
  def initialize(self):
    self.game_id = "%s %d %s" % (datetime.datetime.now().strftime("%Y%m%d %H:%M"), os.getpid(), self.game.team.name[0])
    self.MyLog("*****START GAME*****")
    self.MyLog("my team: %s" % self.game.team.name)
    self.tactics_map = { TZCommander.GO_ENEMY_FLAG:self.GoEnemyFlag, TZCommander.GO_OUR_FLAG:self.GoOurFlag,
                         TZCommander.GO_GOAL:self.GoGoal, TZCommander.GO_ENEMY_GOAL:self.GoEnemyGoal, TZCommander.INTERCEPT_FLAG:self.InterceptFlag,
                         TZCommander.PATROL:self.Patrol, TZCommander.DEFENSE_PATROL:self.DefensePatrol, TZCommander.GO_ENEMY_FLAG_VIA_FLANK:self.GoEnemyFlagViaMyFlank }
    self.verbose = True
    partition = 5
    cols = [i * (self.level.width / partition) for i in range(partition+1)]
    col_centers = []
    for i in range(partition):
      col_centers.append((cols[i] + cols[i+1]) / 2)
    rows = [i * (self.level.height / partition) for i in range(partition+1)]
    row_centers = []
    for i in range(partition):
      row_centers.append((rows[i] + rows[i+1]) / 2)
    self.grid = [[self.level.findNearestFreePosition(Vector2(x, y)) for x in col_centers] for y in row_centers]
    self.flattened_grid = [j for i in self.grid for j in i]
    self.closest_to_our_flag = sorted(self.flattened_grid, key=lambda point: (point - self.game.team.flag.position).squaredLength())

    self.my_alive_pct = 0
    self.opp_alive_pct = 0
    self.my_alive_count =  len(self.game.team.members)
    self.i_have_flag = False
    self.opp_has_flag = False
    self.my_inputs_changed = False
    self.known_match_info_event_length = 0
    
    # maps bot names to last tactic
    self.current_orders = {}
    for bot in self.game.bots_alive:
      self.current_orders[bot.name] = (TZCommander.NONE,{})

  def gen_inputs(self):
    my_alive_count =  len(filter(lambda b: b.health > 0, self.game.team.members))
    my_alive_pct = my_alive_count / float(len(self.game.team.members))
    opp_alive_pct = len(filter(lambda b: b.health > 0 or b.seenlast == None, self.game.enemyTeam.members)) / float(len(self.game.enemyTeam.members))
    opp_has_flag = bool(self.game.team.flag.carrier)
    i_have_flag = bool(self.game.enemyTeam.flag.carrier)

    if (my_alive_count, my_alive_pct, opp_alive_pct, i_have_flag, opp_has_flag) != (self.my_alive_count, self.my_alive_pct, self.opp_alive_pct, self.i_have_flag, self.opp_has_flag):
      self.MyLog("Inputs changed: %s" % ((my_alive_count, my_alive_pct, opp_alive_pct, i_have_flag, opp_has_flag),))
      self.my_alive_count = my_alive_count
      self.my_alive_pct = my_alive_pct
      self.opp_alive_pct = opp_alive_pct
      self.i_have_flag = i_have_flag
      self.opp_has_flag = opp_has_flag
      self.my_inputs_changed = True
    else:
      self.my_inputs_changed = False

    for b in self.game.team.members:
      if b.health <= 0:
        self.current_orders[b.name] = (TZCommander.DEAD, {})
      elif self.current_orders[b.name][0] == TZCommander.DEAD:
        self.current_orders[b.name] = (TZCommander.NONE, {})

  def reassign_roles(self, new_role_count):
    current_role = {}
    # organize team members by tactics
    debug_str = "new_role_count: %s\ncurrent_orders: %s\n" % (new_role_count, self.current_orders)
    for k, v in self.current_orders.iteritems():
      try:
        current_role[v[0]].append(k)
      except KeyError:
        current_role[v[0]] = [k]
    debug_str += "current_role: %s\n" % current_role
    available_bots = []
    for k, v in current_role.iteritems():
      try:
        desired_bots = new_role_count[k][0]
      except KeyError:
        desired_bots = 0
      if k != TZCommander.DEAD:
        for i in range(len(v) - desired_bots):
          available_bots.append(self.game.bots[v.pop()])

        if k!= TZCommander.GO_GOAL:
          # hack to make sure flag carrier is going to goal
          removed_bot_name = None
          for bot_name in v:
            if self.game.bots[bot_name].flag:
              removed_bot_name = bot_name
              self.MyLog("force flag carrier %s to GO_GOAL instead of %s" % (self.game.bots[bot_name], k))
              available_bots.append(self.game.bots[bot_name])
              break
          if removed_bot_name:
            v = filter(lambda b: b != removed_bot_name, v)
            debug_str += "removed flag carrier bot from current_role"
            debug_str += "current_role: %s\n" % current_role
            
    debug_str += "available_bots: %s\n" % available_bots
    if available_bots:
      if len(available_bots) > 1 or removed_bot_name is None:
        self.MyLog("Changing assignment to %s" % new_role_count)
    roles_requested = 0
    new_role_list = list(new_role_count.iteritems())
    # exploit stable sorting to sort by bot count for a role, then by roles that have a function for picking bots
    new_role_list.sort(key=lambda e: len(e[1]) > 1, reverse = True)
    new_role_list.sort(key=lambda e: e[1][0], reverse = True)
    debug_str += "new_role_list: %s\n" % new_role_list
    for k, v in new_role_list:
      try:
        roles_requested += v[0]
        bots_to_add = v[0] - len(current_role[k])
        assert(bots_to_add >= 0)
      except KeyError:
        bots_to_add = v[0]
      if bots_to_add > len(available_bots):
        self.MyLog("ERROR: adding too many bots\nnew_role_count %s\ncurrent_role %s\navaiable_bots %s\n" % (new_role_count, current_role, available_bots), True)
      if bots_to_add > 0 and len(v) > 1:
        available_bots.sort(key=v[1])
        self.MyLog("available_bots sorted by lambda: %s" % available_bots)
      for i in range(bots_to_add):
        self.tactics_map[k](self.game.bots[available_bots.pop(0).name])
    debug_str += "current_orders: %s\n" % self.current_orders
    if roles_requested != self.my_alive_count:
      self.MyLog("ERROR: alive_count: %d roles_requested: %d new_role_count: %s new_role_list: %s" % (
        self.my_alive_count, roles_requested, new_role_count, new_role_list), True)
    assert(roles_requested == self.my_alive_count)
    for b in self.game.team.members:
      if self.current_orders[b.name][0] == TZCommander.NONE:
        self.MyLog("ERROR: bot with NONE orders:\n%s" % debug_str, True)

  def alloc_helper(self, num_bots, alloc):
    new_roles = {}
    if num_bots == 0:
      return new_roles
    bots_remaining = num_bots
    for entry in alloc:
      to_assign = max(1, min(int(round(num_bots * entry[1])), bots_remaining))
      new_roles[entry[0]] = [to_assign]
      new_roles[entry[0]].extend(entry[2:])
      bots_remaining -= to_assign
      if bots_remaining == 0:
        return new_roles
    self.MyLog("ERROR: underassigned %d bots with %s" % (num_bots, alloc), True)
    return new_roles

  def tick(self):
    try:
      self.tick_to_catch()
    except:
      self.MyLog("CAUGHT EXCEPTION:", True)
      self.MyLog(traceback.format_exc(), True)

  def tick_to_catch(self):
    len_events =  len(self.game.match.combatEvents)
    if len_events > self.known_match_info_event_length:
      self.MyLog("New events: %s" % self.game.match.combatEvents[self.known_match_info_event_length:])
      self.known_match_info_event_length = len_events
    self.gen_inputs()
    self.UpdateOrders()
    if self.game.bots_available:
      self.MyLog("bots_available: %s" % self.game.bots_available)
    for b in self.game.bots_available:
      try:
        self.tactics_map[self.current_orders[b.name][0]](b)
      except KeyError:
        self.MyLog("***ERROR key error with %s and current_orders %s" % (b.name, self.current_orders), True)
    for b in filter(lambda b: b.health > 0, self.game.team.members):
      try:
        if self.current_orders[b.name][1]['expires_at'] <= self.game.match.timePassed:
          self.MyLog("reissue due to expiration for %s" % str(self.current_orders[b.name]))
          self.tactics_map[self.current_orders[b.name][0]](b)
      except KeyError:
        pass
      except IndexError:
        self.MyLog("IndexError on %s in %s" % (b.name, self.current_orders), True)

    if self.i_have_flag and self.my_inputs_changed:
      bots_with_flags = filter(lambda b: b.flag, self.game.team.members)
      for b in bots_with_flags:
        if self.current_orders[b.name][0] != TZCommander.GO_GOAL:
          self.MyLog("**bot has flag but not going goal:\n%s\n%s" % (self.current_orders, self.game.team.members))
    if self.my_inputs_changed:
      self.MyLog("current orders: %s" % self.current_orders)

  def MyLog(self, str, also_print = False):
    log_str = "[%s](%.2f) %s" % (self.game_id, self.game.match.timePassed, str)
    self.log.info(log_str)
    if also_print:
      print log_str

  def MyIssue(self, command, bot, pos, description):
    description = str(description)
    self.MyLog("issue %s %s" % (bot.name, description))
    self.issue(command, bot, pos, description = description)

  def GoEnemyFlag(self, bot):
    self.CampLocation(bot, TZCommander.GO_ENEMY_FLAG, self.game.enemyTeam.flag.position)

  def GoOurFlag(self, bot):
    self.CampLocation(bot, TZCommander.GO_OUR_FLAG, self.game.team.flag.position)

  def InterceptFlag(self, bot):
    bot_to_flag = self.game.team.flag.position - bot.position
    if bot_to_flag.squaredLength() == 0:
      tgt = bot.position
    else:
      norm_bot_to_flag = bot_to_flag.normalized()
      interceptor = norm_bot_to_flag.perpendicular()
      half = bot.position.midPoint(self.game.team.flag.position)
      interceptor = interceptor * half.distance(self.game.enemyTeam.flagScoreLocation)
      p = half
      r = interceptor
      q = self.game.team.flag.position
      s = self.game.enemyTeam.flagScoreLocation - self.game.team.flag.position
      u = (q - p).crossProduct(r) / r.crossProduct(s)
      if u<0 or u>1:
        self.MyLog("can't intercept with bot %s flag %s goal %s" % (bot.position, self.game.team.flag.position, self.game.enemyTeam.flagScoreLocation))
        tgt = self.game.enemyTeam.flagScoreLocation
      else:
        tgt = q + s * u
    tgt = self.level.findNearestFreePosition(tgt)
    self.current_orders[bot.name] = (TZCommander.INTERCEPT_FLAG, {'target':tgt, 'expires_at':self.game.match.timePassed + 7})
    if bot.position.squaredDistance(tgt) < 16:
      self.MyLog("%s defending intercept point" % bot.name)
      self.MyIssue(commands.Defend, bot, self.game.team.flag.position - bot.position, self.current_orders[bot.name])
    else:
      self.MyIssue(commands.Attack, bot, tgt, ("%s - mypos %s tgt %s flag %s") % (TZCommander.INTERCEPT_FLAG, bot.position, tgt, self.game.team.flag.position))

  def GoGoal(self, bot):
    self.CampLocation(bot, TZCommander.GO_GOAL, self.game.team.flagScoreLocation)

  def GoEnemyGoal(self, bot):
    self.CampLocation(bot, TZCommander.GO_ENEMY_GOAL, self.game.enemyTeam.flagScoreLocation)
    
  def Patrol(self, bot):
    self.current_orders[bot.name] = (TZCommander.PATROL, {'target':self.PatrolPoints(bot, self.flattened_grid)})

  def DefensePatrol(self, bot):
    self.current_orders[bot.name] = (TZCommander.DEFENSE_PATROL, {'target':self.PatrolPoints(bot, self.closest_to_our_flag[:7], TZCommander.DEFENSE_PATROL)})

  def PatrolPoints(self, bot, waypoints, desc = "Patrol"):
    dest = random.choice(waypoints)
    self.MyIssue(commands.Attack, bot, dest, "%s %s" % (desc, dest))
    return dest

  def AttackBot(self, bot, enemy):
    self.current_orders[bot.name] = (TZCommander.ATTACK_BOT, {'name':enemy.name, 'pos':enemy.position})
    self.MyIssue(commands.Attack, bot, enemy.position, TZCommander.ATTACK_BOT)

  def CampLocation(self, bot, order, target):
    if bot.position.squaredDistance(target) > 4:
      self.MyIssue(commands.Attack, bot, target, order)
      self.current_orders[bot.name] = (order, {'tgt':target})
      return
    direction = random.choice([Vector2.UNIT_X, Vector2.UNIT_Y, Vector2.NEGATIVE_UNIT_X, Vector2.NEGATIVE_UNIT_Y])
    self.current_orders[bot.name] = (order, {'face':direction, 'expires_at':self.game.match.timePassed + random.randint(4,6)})
    self.MyIssue(commands.Defend, bot, direction, self.current_orders[bot.name])

  def GoEnemyFlagViaMyFlank(self, bot):
    self.GoLocViaMyFlank(bot, TZCommander.GO_ENEMY_FLAG_VIA_FLANK, self.game.enemyTeam.flag.position)

  def GoLocViaMyFlank(self, bot, order, target):
    if bot.position.squaredDistance(target) < 100:
      self.CampLocation(bot, order, target)
      return
    diff = target - bot.position
    if abs(diff.x) > abs(diff.y):
      x_dir = 1 if (diff.x > 0) else -1
      waypoint = Vector2(target.x - x_dir * 10, (10 if bot.position.y < target.y else self.level.height - 10))
    else:
      y_dir = 1 if (diff.y > 0) else -1
      waypoint = Vector2((10 if bot.position.x < target.x else self.level.width - 10), target.y - y_dir * 10)
    self.current_orders[bot.name] = (order, {'waypoint':waypoint, 'target':target})
    self.MyIssue(commands.Attack, bot, [self.level.findNearestFreePosition(waypoint), self.level.findNearestFreePosition(target)], order)


class HardCode(TZCommander):
  def UpdateOrders(self):
    if self.i_have_flag:
      new_roles = self.alloc_helper(self.my_alive_count,
                                    ((TZCommander.GO_GOAL, 0.5, lambda b: not b.flag), (TZCommander.DEFENSE_PATROL, 0.5)))
    elif self.opp_has_flag:
      new_roles = self.alloc_helper(self.my_alive_count,
                                    ((TZCommander.INTERCEPT_FLAG, 0.6, lambda b: -b.position.squaredDistance(self.game.team.flag.position)),
                                     (TZCommander.GO_ENEMY_GOAL, 0.4)))
      # new_roles = self.alloc_helper(self.my_alive_count, ((TZCommander.INTERCEPT_FLAG, 1.01),))
    else:
      new_roles = self.alloc_helper(self.my_alive_count,
                                    (
                                      (TZCommander.GO_ENEMY_FLAG_VIA_FLANK, 0.5, lambda b: -b.position.squaredDistance(self.game.enemyTeam.flag.position)),
                                      (TZCommander.GO_OUR_FLAG, 0.25),
                                      (TZCommander.PATROL, 0.25)))
      # new_roles = self.alloc_helper(self.my_alive_count, ((TZCommander.CAMP_FLAG, 1.0),))

    self.reassign_roles(new_roles)

def filename_at_current_path(filename):
  return "%s/%s" % (os.path.dirname(os.path.abspath(inspect.getfile(inspect.currentframe()))), filename)

class NNCommander(TZCommander):
  def initialize(self, nn_dump_file):
    TZCommander.initialize(self)
    self.name = nn_dump_file[0:3]
    self.nn = bpnn.NN(4, 7, len(self.tactics_map))
    nn_info = pickle.load(open(filename_at_current_path(nn_dump_file)))
    self.MyLog("NNCommander with nn_info:\n%s" % nn_info)
    self.nn.wi = nn_info["input_weights"]
    self.nn.wo = nn_info["output_weights"]
    

  def UpdateOrders(self):
    net_results = self.nn.update((self.my_alive_pct, self.opp_alive_pct, float(self.i_have_flag), float(self.opp_has_flag)))
    min_result = min(net_results)
    if min_result < 0:
      net_results = map(lambda x: x - min_result, net_results)
    s = sum(net_results)
    net_results = map(lambda x: x / s, net_results)
    net_results = map(list, zip(self.tactics_map.keys(), net_results))
    for result in net_results:
      if result[0] == TZCommander.GO_GOAL:
        result.append(lambda b: not b.flag)
    net_results.sort(key = lambda x: x[1], reverse = True)
    if self.my_inputs_changed:
      self.MyLog("net_results: %s" % net_results)
    self.reassign_roles(self.alloc_helper(self.my_alive_count, net_results))

class NNCommander1(NNCommander):
  def initialize(self):
    NNCommander.initialize(self, "NN0.dump")

class NNCommander2(NNCommander):
  def initialize(self):
    NNCommander.initialize(self, "NN1.dump")

class NNCommander3(NNCommander):
  def initialize(self):
    NNCommander.initialize(self, "NN2.dump")

class NNCommander4(NNCommander):
  def initialize(self):
    NNCommander.initialize(self, "NN3.dump")
